# coding = utf-8

'''
针对gabor增强的设计
'''

import cv2,os
import numpy as np
import matplotlib.pyplot as plt
import SimpleITK as sitk
from PIL import Image
import math
from skimage import measure

def get_case_id_origion_id_index(case_id, index):
    data_path = "E:\Dataset\Qiye\DongBeiDaXue2\liver"
    i = 0
    origion_id = -1
    for item in sorted(os.listdir(data_path)):
        if case_id - 50 == i:
            origion_id = item
            break
        i += 1
    item_path = os.path.join(data_path, origion_id)
    temp_index = 0
    liver = sitk.GetArrayFromImage(sitk.ReadImage(item_path))
    for i in range(liver.shape[0]):
        if np.max(liver[i]) == 0:
            continue
        if temp_index == index:
            origion_index = i
            break
        temp_index += 1

    return (origion_id, origion_index)




def garbor_filter4(image):
    image = image * 255
    image = image.astype(np.uint8)

    temp = np.zeros(image.shape)

    filters3 = []
    #theta2 = [0, 1 * np.pi / 6, 2 * np.pi / 6, 3 * np.pi / 6, 4 * np.pi / 6, 5 * np.pi / 6]
    theta2 = [0, 1 * np.pi / 6, 2 * np.pi / 6, 3 * np.pi / 6, 4 * np.pi / 6, 5 * np.pi/6]
    for item in theta2:
        kern = cv2.getGaborKernel((2, 2), sigma=1.0, theta=item, lambd=np.pi / 2.0, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        #kern = cv2.getGaborKernel((10, 10), sigma=1.0, theta=item, lambd=5, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters3.append(kern)
    result3 = np.zeros_like(temp)
    for i in range(len(filters3)):
        accum = np.zeros_like(image)
        for kern in filters3[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result3 += np.array(accum)
    result3 = result3 / len(filters3)
    result3 = result3.astype(np.uint8)
    #clahe = cv2.createCLAHE(clipLimit=5, tileGridSize=(100, 100))
    result3 = cv2.equalizeHist(result3)
    #result3 = clahe.apply(result3)

    return result3

def find_data(case_id, origion_id, index):
    big_image = "E:\Dataset\Qiye\DongBeiDaXue2\image_venous\\data2_{}_venous.mha".format(origion_id)
    big_liver = "E:\Dataset\Qiye\DongBeiDaXue2\liver\\data2_{}_liver_label.mha".format(origion_id)
    big_tumor = "E:\Dataset\Qiye\DongBeiDaXue2\lesion\\data2_{}_lesion_label.mha".format(origion_id)
    big_image = sitk.GetArrayFromImage(sitk.ReadImage(big_image))
    big_image[big_image <= -200] = -200
    big_image[big_image > 250] = 250
    big_image = (big_image + 200) / 450
    big_image = big_image[index]
    big_liver = sitk.GetArrayFromImage(sitk.ReadImage(big_liver))
    big_liver = big_liver[index]
    big_tumor = sitk.GetArrayFromImage(sitk.ReadImage(big_tumor))
    big_tumor = big_tumor[index]
    return (big_image, big_liver, big_tumor)


def gabor_standard():

    (origion_id, index) = get_case_id_origion_id_index(case_id=73, index=108)
    origion_id = origion_id.split("_")[1]

    (big_image, big_liver, big_tumor) = find_data(case_id="79", origion_id=origion_id, index=index)

    big_tumor[big_tumor > 0] = 1
    big_tumor = big_tumor * 255
    big_tumor = big_tumor.astype(np.uint8)
    contours, _ = cv2.findContours(big_tumor, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    liver = (big_image) * big_liver
    liver = liver * 255
    liver = liver.astype(np.uint8)
    '''
    liver = cv2.cvtColor(liver, cv2.COLOR_GRAY2BGR)
    for counter in contours:
        data_list = []
        for t in range(counter.shape[0]):
            j = counter[t][0]
            data_list.append(j)
        cv2.polylines(liver, np.array([data_list], np.int32), True, [0, 255, 0], thickness=1)
    '''

    #print(np.max(big_image*big_liver))
    garbor_result = garbor_filter4(big_image * big_liver)
    image = garbor_result
    '''
    image = cv2.cvtColor(garbor_result, cv2.COLOR_GRAY2BGR)
    for counter in contours:
        data_list = []
        for t in range(counter.shape[0]):
            j = counter[t][0]
            data_list.append(j)
        cv2.polylines(image, np.array([data_list], np.int32), True, [0, 255, 0], thickness=1)
    '''

    data = cv2.medianBlur((255 - garbor_result) * big_liver, ksize=5)

    data[data <= 200] = 0
    data[data > 200] = 1

    [data_labels, num] = measure.label(data, return_num=True)
    liver_temp = big_liver.copy()
    liver_temp[liver_temp > 0] = 1
    liver_temp = liver_temp * 255
    liver_temp = liver_temp.astype(np.uint8)
    kernel = np.ones((10, 10), np.uint8)
    erosion = cv2.erode(liver_temp, kernel, iterations=1)
    erosion = liver_temp - erosion
    erosion[erosion > 0] = 1

    for i in np.unique(data_labels):
        temp = np.zeros(data_labels.shape)
        temp[data_labels == i] = 1

        count = float((erosion[temp == 1] == 1).sum()) / float((temp == 1).sum())
        if count >= 0.7:
            data[data_labels == i] = 0



    liver_t = np.zeros(data.shape)
    temp = big_image * big_liver
    liver_t += (temp * data)
    liver_t += (temp * (1 - data) * 1.2)
    liver_t = liver_t / 1.2
    liver_t = liver_t * 255
    liver_t = liver_t.astype(np.uint8)

    '''
    liver_t = cv2.cvtColor(liver_t, cv2.COLOR_GRAY2BGR)
    for counter in contours:
        data_list = []
        for t in range(counter.shape[0]):
            j = counter[t][0]
            data_list.append(j)
        cv2.polylines(liver_t, np.array([data_list], np.int32), True, [0, 255, 0], thickness=1)
    '''

    '''
    plt.subplot(1, 3, 1)
    plt.title("origion image")
    plt.imshow(liver[0:400, 50:450], cmap="gray")
    plt.subplot(1, 3, 2)
    plt.title("gabor enhance")
    plt.imshow(liver_t[0:400, 50:450], cmap="gray")
    plt.subplot(1, 3, 3)
    plt.title("mask")
    plt.imshow(data[0:400, 50:450], cmap="gray")
    plt.show()

    plt.imsave("origion.png", liver[0:400, 50:450], cmap="gray")
    plt.imsave("enhanced.png", liver_t[0:400, 50:450], cmap="gray")
    plt.imsave("enhanced_mask.png", data[0:400, 50:450], cmap="gray")
    plt.imsave("label.png", big_tumor[0:400, 50:450], cmap="gray")
    '''


    data_color = cv2.cvtColor(big_tumor, cv2.COLOR_GRAY2BGR)
    [tumor_labels, num] = measure.label(big_tumor, return_num=True)
    for i in range(data_color.shape[0]):
        for j in range(data_color.shape[1]):
            if tumor_labels[i][j] == 1:
                data_color[i][j] = [138, 43, 226]
            elif tumor_labels[i][j] == 2:
                data_color[i][j] = [218, 112, 214]
            if i >= 222:
                continue
            if data[i][j] > 0 and big_tumor[i][j] == 0:
                data_color[i][j] = [0, 255, 0]



    plt.imsave("multi-label.png", data_color[0:400, 50:450])

#gabor 纹理增强然后设置为蓝色
def gabor_enhanced_blue():
    (origion_id, index) = get_case_id_origion_id_index(case_id=73, index=108)
    origion_id = origion_id.split("_")[1]

    (big_image, big_liver, big_tumor) = find_data(case_id="79", origion_id=origion_id, index=index)

    liver = (big_image) * big_liver
    liver = liver * 255
    liver = liver.astype(np.uint8)

    garbor_result = garbor_filter4(big_image * big_liver)
    image = garbor_result
    data = cv2.medianBlur((255 - garbor_result) * big_liver, ksize=5)
    data[data <= 200] = 0
    data[data > 200] = 1
    [data_labels, num] = measure.label(data, return_num=True)

    data = data * 255
    data = data.astype(np.uint8)
    data = cv2.cvtColor(data, cv2.COLOR_GRAY2BGR)

    data[data_labels == 1] = [47, 79, 79]
    data[data_labels == 2] = [0, 0, 128]
    data[data_labels == 3] = [46, 139, 87]
    data[data_labels == 4] = [34, 139, 34]
    data[data_labels == 5] = [255, 215, 0]
    data[data_labels == 6] = [244, 164, 96]
    data[data_labels == 7] = [255, 69, 0]
    data[data_labels == 8] = [153, 50, 204]
    data[data_labels == 9] = [255, 228, 196]
    data[data_labels == 10] = [122, 103, 238]
    data[data_labels == 11] = [118, 238, 0]
    data[data_labels == 12] = [255, 193, 37]
    data[data_labels == 13] = [238, 18, 137]
    data[data_labels == 14] = [155, 48, 255]
    data[data_labels == 15] = [144, 238, 144]

    liver_temp = big_liver.copy()
    liver_temp[liver_temp > 0] = 1
    liver_temp = liver_temp * 255
    liver_temp = liver_temp.astype(np.uint8)
    kernel = np.ones((10, 10), np.uint8)
    erosion = cv2.erode(liver_temp, kernel, iterations=1)
    erosion = liver_temp - erosion
    erosion[erosion > 0] = 1

    for i in np.unique(data_labels):
        temp = np.zeros(data_labels.shape)
        temp[data_labels == i] = 1

        count = float((erosion[temp == 1] == 1).sum()) / float((temp == 1).sum())
        if count >= 0.7:
            data[data_labels == i] = [0, 0, 0]

    plt.imsave("sequence/remove_erison.png", data[0:400, 50:450])







if __name__ == '__main__':
    gabor_enhanced_blue()